from langgraph_coder.src.agent.model.graph_state import GraphState


class FinishDispatch:
    def __init__(self):
        pass

    def __call__(self, state: GraphState) -> GraphState:
        """
        完成分发
        :param state:
        :return:
        """
        error = state["error"]
        iterations = state["iterations"]
        max_iterations = state["max_iterations"]
        flag = state["flag"]
        if error == "no" or iterations == max_iterations:
            print("---决定完成---")
            return "end"
        else:
            print("---决定继续生成代码---")
            if flag == "reflect":
                return "reflect"
            else:
                return "generate"
